{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Efficiently scale distributed training Llama 3, PyTorch FSDP and Q-Lora\n",
    "\n",
    "Open LLMs like Meta [Llama 3](https://huggingface.co/meta-llama/Meta-Llama-3-70b), Mistral AI [Mistral](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2) & [Mixtral](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) models or AI21 [Jamba](https://huggingface.co/ai21labs/Jamba-v0.1) are now OpenAI competitors. However, most of the time you need to fine-tune the model on your data to unlock the full potential of the model. Fine-tuning smaller LLMs, like Mistral became very accessible on a single GPU by using Q-Lora. But efficiently fine-tuning bigger models like Llama 3 70b or Mixtral stayed a challenge until now. \n",
    "\n",
    "This blog post walks you thorugh how to fine-tune a Llama 3 using PyTorch FSDP and Q-Lora with the help of Hugging Face [TRL](https://huggingface.co/docs/trl/index), [Transformers](https://huggingface.co/docs/transformers/index), [peft](https://huggingface.co/docs/peft/index) & [datasets](https://huggingface.co/docs/datasets/index). In addition to FSDP we will use [Flash Attention v2 through the Pytorch SDPA](https://pytorch.org/blog/pytorch2-2/) implementation. \n",
    "\n",
    "1. [Setup development environment](#1-setup-development-environment)\n",
    "2. [Create and prepare the dataset](#2-create-and-prepare-the-dataset)\n",
    "3. [Fine-tune the LLM with PyTorch FSDP, Q-Lora and SDPA](#3-fine-tune-the-llm-with-pytorch-fsdp-q-lora-and-sdpa)\n",
    "4. [Test Model and run Inference](#4-test-model-and-run-inference)\n",
    "\n",
    "_Note: This blog was created and validated on NVIDIA H100 and NVIDIA A10G GPUs. The configurations and code is optimized for 4xA10G GPUs each with 24GB of Memory. I hope this keeps the example as accessible as possible for most people. If you have access to more compute you can make changes to the config (yaml) at step 3._\n",
    "\n",
    "**FSDP + Q-Lora Background**\n",
    "\n",
    "In a collaboration between [Answer.AI](https://www.answer.ai/posts/2024-03-06-fsdp-qlora.html), Tim Dettmers [Q-Lora creator](https://github.com/TimDettmers/bitsandbytes) and [Hugging Face](https://huggingface.co/), we are proud to announce to share the support of Q-Lora and PyTorch FSDP (Fully Sharded Data Parallel). FSDP and Q-Lora allows you now to fine-tune Llama 2 70b or Mixtral 8x7B on 2x consumer GPUs (24GB). If you want to learn more about the background of this collaboration take a look at [You can now train a 70b language model at home](https://www.answer.ai/posts/2024-03-06-fsdp-qlora.html). Hugging Face PEFT is were the magic happens for this happens, read more about it in the [PEFT documentation](https://huggingface.co/docs/peft/v0.10.0/en/accelerate/fsdp).\n",
    "\n",
    "* [PyTorch FSDP](https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/) is a data/model parallelism technique that shards model across GPUs, reducing memory requirements and enabling the training of larger models more efficiently​​​​​​.\n",
    "* Q-LoRA is a fine-tuning method that leverages quantization and Low-Rank Adapters to efficiently reduced computational requirements and memory footprint. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup development environment\n",
    "\n",
    "Our first step is to install Hugging Face Libraries and Pyroch, including trl, transformers and datasets. If you haven't heard of trl yet, don't worry. It is a new library on top of transformers and datasets, which makes it easier to fine-tune, rlhf, align open LLMs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install Pytorch for FSDP and FA/SDPA\n",
    "%pip install \"torch==2.2.2\" tensorboard\n",
    "\n",
    "# Install Hugging Face libraries\n",
    "%pip install  --upgrade \"transformers==4.40.0\" \"datasets==2.18.0\" \"accelerate==0.29.3\" \"evaluate==0.4.1\" \"bitsandbytes==0.43.1\" \"huggingface_hub==0.22.2\" \"trl==0.8.6\" \"peft==0.10.0\" "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we need to login into Hugging Face to access the Llama 3 70b model. If you don't have an account yet and accepted the terms, you can create one [here](https://huggingface.co/join). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!huggingface-cli login --token \"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Create and prepare the dataset\n",
    "\n",
    "After our environment is set up, we can start creating and preparing our dataset. A fine-tuning dataset should have a diverse set of demonstrations of the task you want to solve. If you want to learn more about how to create a dataset, take a look at the [How to Fine-Tune LLMs in 2024 with Hugging Face](https://www.philschmid.de/fine-tune-llms-in-2024-with-trl#3-create-and-prepare-the-dataset).\n",
    "\n",
    "We will use the [HuggingFaceH4/no_robots](https://huggingface.co/datasets/HuggingFaceH4/no_robots) dataset a high-quality dataset of 10,000 instructions and demonstrations created by skilled human annotators. This data can be used for supervised fine-tuning (SFT) to make language models follow instructions better. No Robots was modelled after the instruction dataset described in OpenAI's [InstructGPT paper](https://huggingface.co/papers/2203.02155), and is comprised mostly of single-turn instructions.\n",
    "\n",
    "```json\n",
    "{\"messages\": [{\"role\": \"system\", \"content\": \"You are...\"}, {\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"...\"}]}\n",
    "{\"messages\": [{\"role\": \"system\", \"content\": \"You are...\"}, {\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"...\"}]}\n",
    "{\"messages\": [{\"role\": \"system\", \"content\": \"You are...\"}, {\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"...\"}]}\n",
    "```\n",
    "\n",
    "The [no_robots](https://huggingface.co/datasets/HuggingFaceH4/no_robots) dataset has 10,000 split into 9,500 training and  500 test examples. Some samples are not including a `system` message. We will load the dataset with the `datasets` library, add a missing `system` message and save them to separate json files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "# Convert dataset to OAI messages\n",
    "system_message = \"\"\"You are Llama, an AI assistant created by Philipp to be helpful and honest. Your knowledge spans a wide range of topics, allowing you to engage in substantive conversations and provide analysis on complex subjects.\"\"\"\n",
    "\n",
    "def create_conversation(sample):\n",
    "    if sample[\"messages\"][0][\"role\"] == \"system\":\n",
    "        return sample\n",
    "    else:\n",
    "      sample[\"messages\"] = [{\"role\": \"system\", \"content\": system_message}] + sample[\"messages\"]\n",
    "      return sample\n",
    "\n",
    "# Load dataset from the hub\n",
    "dataset = load_dataset(\"HuggingFaceH4/no_robots\")\n",
    "\n",
    "# Add system message to each conversation\n",
    "columns_to_remove = list(dataset[\"train\"].features)\n",
    "columns_to_remove.remove(\"messages\")\n",
    "dataset = dataset.map(create_conversation, remove_columns=columns_to_remove,batched=False)\n",
    "\n",
    "# Filter out conversations which are corrupted with wrong turns, keep which have even number of turns after adding system message\n",
    "dataset[\"train\"] = dataset[\"train\"].filter(lambda x: len(x[\"messages\"][1:]) % 2 == 0)\n",
    "dataset[\"test\"] = dataset[\"test\"].filter(lambda x: len(x[\"messages\"][1:]) % 2 == 0)\n",
    "\n",
    "# save datasets to disk \n",
    "dataset[\"train\"].to_json(\"train_dataset.json\", orient=\"records\", force_ascii=False)\n",
    "dataset[\"test\"].to_json(\"test_dataset.json\", orient=\"records\", force_ascii=False)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Fine-tune the LLM with PyTorch FSDP, Q-Lora and SDPA\n",
    "\n",
    "We are now ready to fine-tune our model with PyTorch FSDP, Q-Lora and SDPA. Since we are running in a distributed setup, we need to use `torchrun` and a python script to start the training. \n",
    "\n",
    "We prepared a script [run_fsdp_qlora.py](./scripts/run_fsdp_qlora.py) which will load the dataset from disk, prepare the model, tokenizer and start the training. It usees the [SFTTrainer](https://huggingface.co/docs/trl/sft_trainer) from `trl` to fine-tune our model. The `SFTTrainer` makes it straightfoward to supervise fine-tune open LLMs supporting:\n",
    "* Dataset formatting, including conversational and instruction format (✅ used)\n",
    "* Training on completions only, ignoring prompts (❌ not used)\n",
    "* Packing datasets for more efficient training (✅ used)\n",
    "* PEFT (parameter-efficient fine-tuning) support including Q-LoRA (✅ used)\n",
    "* Preparing the model and tokenizer for conversational fine-tuning (❌ not used, see below)\n",
    "\n",
    "_Note: We are using an Anthropic/Vicuna like Chat Template with `User:` and `Assistant:` roles. This done because the special tokens in base Llama 3 (`<|begin_of_text|>` or `<|reserved_special_token_XX|>`) are not trained. Meaning if want would like to use them for the template we need to train them which requires more memory, since we need to update the embedding layer and lm_head. If you have access to more compute you can modify `LLAMA_3_CHAT_TEMPLATE` in the [run_fsdp_qlora.py](./scripts/run_fsdp_qlora.py) script._\n",
    "\n",
    "For configuration we use the new `TrlParser`, that allows us to provide hyperparameters in a yaml file or overwrite the arguments from the config file by explicitly passing them to the CLI, e.g. `--num_epochs 10`. Below is the config file for fine-tuning Llama 3 70B on 4x A10G GPUs or 4x24GB GPUs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile llama_3_70b_fsdp_qlora.yaml\n",
    "# script parameters\n",
    "model_id: \"meta-llama/Meta-Llama-3-70b\" # Hugging Face model id\n",
    "dataset_path: \".\"                      # path to dataset\n",
    "max_seq_len:  3072 # 2048              # max sequence length for model and packing of the dataset\n",
    "# training parameters\n",
    "output_dir: \"./llama-3-70b-hf-no-robot\" # Temporary output directory for model checkpoints\n",
    "report_to: \"tensorboard\"               # report metrics to tensorboard\n",
    "learning_rate: 0.0002                  # learning rate 2e-4\n",
    "lr_scheduler_type: \"constant\"          # learning rate scheduler\n",
    "num_train_epochs: 3                    # number of training epochs\n",
    "per_device_train_batch_size: 1         # batch size per device during training\n",
    "per_device_eval_batch_size: 1          # batch size for evaluation\n",
    "gradient_accumulation_steps: 2         # number of steps before performing a backward/update pass\n",
    "optim: adamw_torch                     # use torch adamw optimizer\n",
    "logging_steps: 10                      # log every 10 steps\n",
    "save_strategy: epoch                   # save checkpoint every epoch\n",
    "evaluation_strategy: epoch             # evaluate every epoch\n",
    "max_grad_norm: 0.3                     # max gradient norm\n",
    "warmup_ratio: 0.03                     # warmup ratio\n",
    "bf16: true                             # use bfloat16 precision\n",
    "tf32: true                             # use tf32 precision\n",
    "gradient_checkpointing: true           # use gradient checkpointing to save memory\n",
    "# FSDP parameters: https://huggingface.co/docs/transformers/main/en/fsdp\n",
    "fsdp: \"full_shard auto_wrap offload\" # remove offload if enough GPU memory\n",
    "fsdp_config:\n",
    "  backward_prefetch: \"backward_pre\"\n",
    "  forward_prefetch: \"false\"\n",
    "  use_orig_params: \"false\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_Note: At the end of the training there will be a slight increase in GPU memory usage (~10%). This is due to the saving of the model correctly. Make sure to have enough memory left on your GPU to save the model. [REF](https://huggingface.co/docs/peft/v0.10.0/en/accelerate/fsdp#memory-usage)_\n",
    "\n",
    "To launch our training we will use `torchrun` to keep the example flexible and easy to adjust to, e.g. Amazon SageMaker or Google Cloud Vertex AI. For `torchrun` and FSDP we need to set the environment variable `ACCELERATE_USE_FSDP` and `FSDP_CPU_RAM_EFFICIENT_LOADING` to tell transformers/accelerate to use FSDP and load the model in a memory-efficient way.  \n",
    "\n",
    "_Note: To NOT CPU offloading you need to change the value of `fsdp` and remove `offload`. This only works on > 40GB GPUs since it requires more memory._\n",
    "\n",
    "Now, lets launch the training with the following command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ACCELERATE_USE_FSDP=1 FSDP_CPU_RAM_EFFICIENT_LOADING=1 torchrun --nproc_per_node=4 ./scripts/run_fsdp_qlora.py --config llama_3_70b_fsdp_qlora.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Memory usage:**\n",
    "* Full-finetuning with FSDP needs ~16X80GB GPUs\n",
    "* FSDP + LoRA needs ~8X80GB GPUs\n",
    "* FSDP + Q-Lora needs ~2x40GB GPUs \n",
    "* FSDP + Q-Lora + CPU offloading needs 4x24GB GPUs, with 22 GB/GPU and 127 GB CPU RAM with a sequence length of 3072 and a batch size of 1. \n",
    "\n",
    "The training of Llama 3 70B with Flash Attention for 3 epochs with a dataset of 10k samples takes 45h on a `g5.12xlarge`. The instance costs `5.67$/h` which would result in a total cost of `255.15$`. This sounds expensive but allows you to fine-tune a Llama 3 70B on small GPU resources. If we scale up the training to 4x H100 GPUs, the training time will be reduced to ~1,25h. If we assume 1x H100 costs `5-10$/h` the total cost would between `25$-50$`.\n",
    "\n",
    "We can see a trade-off between accessibility and performance. If you have access to more/better compute you can reduce the training time and cost, but even with small resources you can fine-tune a Llama 3 70B. The cost/performance is different since for 4x A10G GPUs we need to offload the model to the CPU which reduces the overall flops. \n",
    "\n",
    "_Note: During evaluation and testing of the blog post I noticed that ~40 max steps (80 samples stacked to 3k sequence length) are enough for first results. The training for 40 steps ~1h or ~$5._"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### _Optional: Merge LoRA adapter in to the original model_\n",
    "\n",
    "When using QLoRA, we only train adapters and not the full model. This means when saving the model during training we only save the adapter weights and not the full model. If you want to save the full model, which makes it easier to use with Text Generation Inference you can merge the adapter weights into the model weights using the `merge_and_unload` method and then save the model with the `save_pretrained` method. This will save a default model, which can be used for inference.\n",
    "\n",
    "_Note: You might require > 192GB CPU Memory._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### COMMENT IN TO MERGE PEFT AND BASE MODEL ####\n",
    "# from peft import AutoPeftModelForCausalLM\n",
    "\n",
    "# # Load PEFT model on CPU\n",
    "# model = AutoPeftModelForCausalLM.from_pretrained(\n",
    "#     args.output_dir,\n",
    "#     torch_dtype=torch.float16,\n",
    "#     low_cpu_mem_usage=True,\n",
    "# )  \n",
    "# # Merge LoRA and base model and save\n",
    "# merged_model = model.merge_and_unload()\n",
    "# merged_model.save_pretrained(args.output_dir,safe_serialization=True, max_shard_size=\"2GB\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Test Model and run Inference\n",
    "\n",
    "After the training is done we want to evaluate and test our model. We will load different samples from the original dataset and evaluate the model manually. Evaluating Generative AI models is not a trivial task since 1 input can have multiple correct outputs. If you want to learn more about evaluating generative models, check out [Evaluate LLMs and RAG a practical example using Langchain and Hugging Face](https://www.philschmid.de/evaluate-llm) blog post."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from peft import AutoPeftModelForCausalLM\n",
    "from transformers import AutoTokenizer \n",
    "\n",
    "peft_model_id = \"./llama-3-70b-hf-no-robot\"\n",
    "\n",
    "# Load Model with PEFT adapter\n",
    "model = AutoPeftModelForCausalLM.from_pretrained(\n",
    "  peft_model_id,\n",
    "  torch_dtype=torch.float16,\n",
    "  quantization_config= {\"load_in_4bit\": True},\n",
    "  device_map=\"auto\"\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(peft_model_id)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s load our test dataset try to generate an instruction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset \n",
    "from random import randint\n",
    "\n",
    "\n",
    "# Load our test dataset\n",
    "eval_dataset = load_dataset(\"json\", data_files=\"test_dataset.json\", split=\"train\")\n",
    "rand_idx = randint(0, len(eval_dataset))\n",
    "messages = eval_dataset[rand_idx][\"messages\"][:2]\n",
    "\n",
    "# Test on sample \n",
    "input_ids = tokenizer.apply_chat_template(messages,add_generation_prompt=True,return_tensors=\"pt\").to(model.device)\n",
    "outputs = model.generate(\n",
    "    input_ids,\n",
    "    max_new_tokens=512,\n",
    "    eos_token_id= tokenizer.eos_token_id,\n",
    "    do_sample=True,\n",
    "    temperature=0.6,\n",
    "    top_p=0.9,\n",
    ")\n",
    "response = outputs[0][input_ids.shape[-1]:]\n",
    "\n",
    "print(f\"**Query:**\\n{eval_dataset[rand_idx]['messages'][1]['content']}\\n\")\n",
    "print(f\"**Original Answer:**\\n{eval_dataset[rand_idx]['messages'][2]['content']}\\n\")\n",
    "print(f\"**Generated Answer:**\\n{tokenizer.decode(response,skip_special_tokens=True)}\")\n",
    "\n",
    "# **Query:**\n",
    "# How long was the Revolutionary War?\n",
    "# **Original Answer:**\n",
    "# The American Revolutionary War lasted just over seven years. The war started on April 19, 1775, and ended on September 3, 1783. \n",
    "# **Generated Answer:**\n",
    "# The Revolutionary War, also known as the American Revolution, was an 18th-century war fought between the Kingdom of Great Britain and the Thirteen Colonies. The war lasted from 1775 to 1783."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That looks pretty good! 🚀 Now, its your turn! \n",
    "\n",
    "If you want to deploy your model into production check out [Deploy the LLM for Production](https://www.philschmid.de/fine-tune-llms-in-2024-with-trl#6-deploy-the-llm-for-production)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2d58e898dde0263bc564c6968b04150abacfd33eed9b19aaa8e45c040360e146"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
